# -*- coding: utf-8 -*-
# @Time    : 2023-12-14 23:24
# @Author  : Jiang Liu

import cv2


def show_xml_label(img_path, xml_path):
    img = cv2.imread(img_path)
    import xml.etree.ElementTree as ET
    tree = ET.parse(xml_path)
    root = tree.getroot()
    for obj in root.findall('object'):
        name = obj.find('name').text
        bndbox = obj.find('bndbox')
        xmin = int(bndbox.find('xmin').text)
        ymin = int(bndbox.find('ymin').text)
        xmax = int(bndbox.find('xmax').text)
        ymax = int(bndbox.find('ymax').text)
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
        # cv2.putText(img, name, (xmin, ymin - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
    return img


def show_yolo_label(img_path, txt_path):
    img = cv2.imread(img_path)
    with open(txt_path, 'r') as f:
        for line in f.readlines():
            line = line.strip().split(' ')
            cls_id = int(line[0])
            x = float(line[1])
            y = float(line[2])
            w = float(line[3])
            h = float(line[4])
            xmin = int((x - w / 2) * img.shape[1])
            ymin = int((y - h / 2) * img.shape[0])
            xmax = int((x + w / 2) * img.shape[1])
            ymax = int((y + h / 2) * img.shape[0])
            cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
    return img


def main():
    # basename = '0000521 (403)'
    basename = '0000521 (400)'
    img_path = f'VOC2023/JPEGImages/{basename}.png'
    xml_path = f'VOC2023/Annotations/{basename}.xml'
    # txt_path = f'../dataset_demo/txts/{basename}.txt'

    img_original = cv2.imread(img_path)
    img_xml_label = show_xml_label(img_path, xml_path)
    # img_txt_label = show_yolo_label(img_path, txt_path)

    cv2.imwrite(f'test_out/{basename}_original.png', img_original)
    cv2.imwrite(f'test_out/{basename}_xml_label.png', img_xml_label)
    # cv2.imwrite(f'../test_out/{basename}_txt_label.png', img_txt_label)


if __name__ == '__main__':
    main()
